package org.xi.maple.sink

import org.apache.commons.lang3.StringUtils
import org.apache.spark.sql.execution.datasources.jdbc.{JDBCOptions, JdbcUtils}
import org.apache.spark.sql.{Dataset, Row, SparkSession}
import org.slf4j.{Logger, LoggerFactory}
import org.xi.maple.api.MapleSink

import java.sql.Connection
import scala.collection.JavaConverters._
import scala.collection.mutable

class Jdbc extends MapleSink[JdbcConfig] {

  private val log: Logger = LoggerFactory.getLogger(classOf[Jdbc])

  override def checkConfig: (Boolean, String) = {
    val nonExistsOptions = mutable.MutableList[String]()
    if (StringUtils.isBlank(config.getUrl)) {
      nonExistsOptions.+=:("url")
    }
    if (StringUtils.isBlank(config.getDriver)) {
      nonExistsOptions.+=:("driver")
    }
    if (StringUtils.isBlank(config.getTargetTable)) {
      nonExistsOptions.+=:("targetTable")
    }
    if (StringUtils.isBlank(config.getUser)) {
      nonExistsOptions.+=:("username")
    }

    if (nonExistsOptions.isEmpty) {
      (true, "")
    } else {
      (false, "please specify " + nonExistsOptions.map { field => s"[${field}]" }.mkString(", ") + " as non-empty string")
    }
  }

  def output(spark: SparkSession, argsMap: mutable.Map[String, String], ds: Dataset[Row]): Unit = {

    var options = Map(
      "url" -> config.getUrl,
      "driver" -> config.getDriver,
      "user" -> config.getUser,
      "password" -> config.getPassword,
      "dbtable" -> config.getTargetTable
    )

    if (config.getOptions != null && !config.getOptions.isEmpty) {
      options = config.getOptions.asScala.toMap ++ options
    }

    if (config.getPreQueries != null && !config.getPreQueries.isEmpty) {
      config.getPreQueries.asScala.foreach(query => {
        log.info(s"Execute pre query: $query")
        execute(options, query)
      })
    }

    val writer = if (config.getNumPartitions != null) {
      ds.repartition(config.getNumPartitions).write.format("jdbc")
    } else {
      ds.write.format("jdbc")
    }
    if (StringUtils.isNotBlank(config.getSaveMode)) {
      writer.mode(config.getSaveMode)
    }
    log.info(s"Save data to jdbc url: ${config.getUrl}, driver: ${config.getDriver}, username: ${config.getUser}, table: ${config.getTargetTable}")
    writer.options(options).save()
  }

  private def execute(options: Map[String, String], query: String): Unit = {
    log.info("Execute query: {}", query)
    val jdbcOptions = new JDBCOptions(options)
    val conn: Connection = JdbcUtils.createConnectionFactory(jdbcOptions)()
    try {
      val statement = conn.prepareStatement(query)
      try {
        statement.setQueryTimeout(jdbcOptions.queryTimeout)
        val rows = statement.executeUpdate()
        log.info("{} rows affected", rows)
      } finally {
        statement.close()
      }
    } catch {
      case e: Exception => log.error("Execute query failed. ", e)
    } finally {
      conn.close()
    }
  }
}
